import os
import qianfan

# This automatically reads the keys from your environment variables!
os.environ["QIANFAN_ACCESS_KEY"]=os.environ["QIANFAN_ACCESS_KEY"]
os.environ["QIANFAN_SECRET_KEY"]=os.environ["QIANFAN_SECRET_KEY"]

def generate_poem(topic):
    """
    Uses the Wenxin LLM to generate a short poem about a given topic.
    """
    chat_comp = qianfan.ChatCompletion()

    # This is our request to the "waiter" (API)
    resp = chat_comp.do(
        model="ERNIE-Bot-4",
        messages=[{
            "role": "user",
            "content": f"Please write a short, four-line poem about {topic}."
        }]
    )

    # The response from the "kitchen" (LLM) is inside the result
    poem = resp["result"]
    return poem

print("Welcome to the AI Poem Generator!")
poem_topic = input("What should the poem be about?")

ai_poem = generate_poem(poem_topic)

print("\nHere is your poem:")
print("--------------------")
print(ai_poem)

